import cv2
import os
import json
import numpy as np

# 文件夹路径
image_dir = "/home/champrin/record_data/网盘/能量机关/24-国赛-场地道具训练/extract"  # 图片文件夹路径
label_dir = "/home/champrin/record_data/网盘/能量机关/24-国赛-场地道具训练/extract"  # JSON标注文件夹路径
output_dir = "/home/champrin/record_data/网盘/能量机关/24-国赛-场地道具训练/extract/output"  # 输出图像文件夹路径

os.makedirs(output_dir, exist_ok=True)

# 遍历图片文件夹
for img_name in os.listdir(image_dir):
    if not img_name.lower().endswith(('.jpg', '.png', '.jpeg')):
        continue

    img_path = os.path.join(image_dir, img_name)
    label_path = os.path.join(label_dir, os.path.splitext(img_name)[0] + ".json")

    if not os.path.exists(label_path):
        print(f"跳过 {img_name}，未找到对应标注文件")
        continue

    # 读取图片
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图片: {img_name}")
        continue

    # 读取标注文件
    with open(label_path, 'r') as f:
        data = json.load(f)

    # 遍历每个对象的分割区域并填充
    for obj in data['objects']:
        segment = obj.get("segmentation", [])

        points = np.array(segment, dtype=np.int32).reshape((-1, 1, 2))
        cv2.fillPoly(img, [points], (0, 0, 0))  # 画黑色区域

    # 保存处理后的图像
    output_img_path = os.path.join(output_dir, img_name)
    cv2.imwrite(output_img_path, img)
    print(f"已保存: {output_img_path}")

# 完成
print("所有图像处理完成！")
